# -*- coding: utf-8 -*-

import os
import json
import logging
from typing import Dict, Any, List, Union
import numpy as np

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class ResultManager:
    """
    结果管理器类，负责保存和管理实验结果
    """
    
    def __init__(self, config):
        """
        初始化结果管理器
        
        Args:
            config: 配置对象
        """
        self.config = config
        # 确保输出目录存在
        self._ensure_output_directory_exists()
        
    
    def _ensure_output_directory_exists(self):
        """
        确保输出目录存在，如果不存在则创建
        """
        output_dir = getattr(self.config, 'OUTPUT_DIR', 'results')
        if not os.path.exists(output_dir):
            try:
                os.makedirs(output_dir)
                logger.info(f"创建输出目录: {output_dir}")
            except Exception as e:
                logger.error(f"创建输出目录失败: {str(e)}")
    
    
    def save_results(self, results: List[Dict[str, Any]], output_file: str) -> bool:
        """
        保存实验结果到JSON文件
        
        Args:
            results: 实验结果列表或字典
            output_file: 输出文件路径
            
        Returns:
            保存是否成功
        """
        try:
            # 确保输出目录存在
            output_dir = os.path.dirname(output_file)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)
                
            # 处理可能的不可序列化对象
            serializable_results = self._convert_to_serializable(results)
            
            # 保存结果到文件
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(serializable_results, f, ensure_ascii=False, indent=2)
                
            logger.info(f"实验结果已成功保存到: {output_file}")
            return True
            
        except Exception as e:
            logger.error(f"保存实验结果失败: {str(e)}")
            import traceback
            logger.error(f"详细错误栈: {traceback.format_exc()}")
            return False
    
    
    def append_result(
        self,
        result: Dict[str, Any],
        output_file: str,
        experiment_info: Dict[str, Any]
    ) -> bool:
        """
        追加单个实验结果到文件
        
        Args:
            result: 单个实验结果
            output_file: 输出文件路径
            experiment_info: 实验信息
            
        Returns:
            追加是否成功
        """
        try:
            # 确保输出目录存在
            output_dir = os.path.dirname(output_file)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)
            
            # 检查文件是否存在
            existing_results = []
            if os.path.exists(output_file):
                try:
                    with open(output_file, 'r', encoding='utf-8') as f:
                        existing_results = json.load(f)
                    # 确保是列表格式
                    if not isinstance(existing_results, list):
                        existing_results = [existing_results]
                except json.JSONDecodeError:
                    # 文件存在但内容无效，重新开始
                    existing_results = []
            
            # 将结果转换为可序列化格式
            serializable_result = self._convert_to_serializable(result)
            
            # 查找对应的实验组或创建新组
            experiment_id = experiment_info.get('experiment_id')
            experiment_updated = False
            
            for exp in existing_results:
                if exp.get('experiment_id') == experiment_id:
                    # 找到对应的实验组，追加结果
                    if 'results' not in exp:
                        exp['results'] = []
                    exp['results'].append(serializable_result)
                    exp['total_samples'] = len(exp['results'])
                    # 更新平均值
                    exp['average_metrics'] = self._calculate_average_metrics(exp['results'])
                    experiment_updated = True
                    break
            
            if not experiment_updated:
                # 创建新的实验组
                new_experiment = experiment_info.copy()
                new_experiment['results'] = [serializable_result]
                new_experiment['total_samples'] = 1
                new_experiment['average_metrics'] = self._calculate_average_metrics(new_experiment['results'])
                existing_results.append(new_experiment)
            
            # 保存更新后的结果
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(existing_results, f, ensure_ascii=False, indent=2)
            
            return True
            
        except Exception as e:
            logger.error(f"追加实验结果失败: {str(e)}")
            import traceback
            logger.error(f"详细错误栈: {traceback.format_exc()}")
            return False
    
    
    def _convert_to_serializable(self, obj: Any) -> Any:
        """
        递归将对象转换为可JSON序列化的格式
        
        Args:
            obj: 任意对象
            
        Returns:
            可序列化的对象
        """
        # 处理字典
        if isinstance(obj, dict):
            return {key: self._convert_to_serializable(value) for key, value in obj.items()}
        
        # 处理列表
        elif isinstance(obj, list):
            return [self._convert_to_serializable(item) for item in obj]
        
        # 处理numpy数值类型
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, np.generic):
            return obj.item()
        
        # 处理PyTorch张量和其他可能的数值类型
        elif hasattr(obj, 'item'):
            try:
                return obj.item()
            except:
                # 如果转换失败，转为字符串
                return str(obj)
        
        # 基本数据类型可以直接序列化
        elif isinstance(obj, (int, float, str, bool, type(None))):
            return obj
        
        # 其他类型转换为字符串
        else:
            return str(obj)
    
    
    def _calculate_average_metrics(self, results: List[Dict[str, Any]]) -> Dict[str, float]:
        """
        计算实验结果的平均指标
        
        Args:
            results: 实验结果列表
            
        Returns:
            平均指标字典
        """
        if not results:
            return {}
        
        # 收集所有数值指标
        all_metrics = {}
        for result in results:
            # 跳过错误记录
            if 'error' in result:
                continue
            
            # 递归收集嵌套对象中的数值指标
            self._collect_metrics(result, all_metrics, prefix='')
        
        # 计算平均值
        average_metrics = {}
        for metric_name, values in all_metrics.items():
            if values:
                average_value = sum(values) / len(values)
                # 保留四位小数
                average_metrics[metric_name] = round(average_value, 4)
        
        return average_metrics
    
    
    def _collect_metrics(self, obj: Any, metrics_dict: Dict[str, List[float]], prefix: str = '') -> None:
        """
        递归收集嵌套对象中的数值指标
        
        Args:
            obj: 任意对象
            metrics_dict: 用于收集指标的字典
            prefix: 指标名称前缀
        """
        if isinstance(obj, dict):
            for key, value in obj.items():
                # 为嵌套键添加前缀，避免末尾点号问题
                if prefix:
                    new_prefix = f"{prefix}.{key}"
                else:
                    new_prefix = key
                self._collect_metrics(value, metrics_dict, new_prefix)
        elif isinstance(obj, list):
            for i, item in enumerate(obj):
                new_prefix = f"{prefix}[{i}]"
                self._collect_metrics(item, metrics_dict, new_prefix)
        elif isinstance(obj, (int, float)):
            # 直接使用当前prefix作为指标名称
            metric_name = prefix
            if metric_name not in metrics_dict:
                metrics_dict[metric_name] = []
            metrics_dict[metric_name].append(obj)


# 为了兼容旧的函数式调用方式，提供全局函数

def save_results(results, output_file):
    """兼容旧版本的函数式调用"""
    # 创建一个简单的配置对象
    class DummyConfig:
        def __init__(self):
            self.OUTPUT_DIR = os.path.dirname(output_file) if os.path.dirname(output_file) else 'results'
    
    config = DummyConfig()
    result_manager = ResultManager(config)
    return result_manager.save_results(results, output_file)


def append_result(result, output_file, experiment_info=None):
    """兼容旧版本的函数式调用"""
    # 创建一个简单的配置对象
    class DummyConfig:
        def __init__(self):
            self.OUTPUT_DIR = os.path.dirname(output_file) if os.path.dirname(output_file) else 'results'
    
    config = DummyConfig()
    result_manager = ResultManager(config)
    
    # 如果没有提供实验信息，创建一个默认的
    if experiment_info is None:
        from datetime import datetime
        experiment_info = {
            "experiment_id": "default_experiment",
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
    
    return result_manager.append_result(result, output_file, experiment_info)